"""
图片审核算法
"""
import os
import requests
from typing import List, Dict
from PIL import Image
import numpy as np
from utils.logger import audit_logger
from models.schemas import ImageAuditResult


class ImageAuditor:
    """图片审核器"""
    
    def __init__(self):
        self.violation_keywords = {
            'adult': ['成人内容', '色情', '裸体'],
            'violence': ['暴力', '血腥', '武器'],
            'illegal': ['毒品', '赌博', '违法'],
            'spam': ['广告', '二维码', '推广']
        }
    
    def audit_image(self, image_url: str) -> ImageAuditResult:
        """审核单张图片"""
        try:
            if not image_url:
                return ImageAuditResult(
                    is_safe=True,
                    violations=[],
                    confidence=1.0
                )
            
            # 下载图片
            image = self._download_image(image_url)
            if image is None:
                return ImageAuditResult(
                    is_safe=True,
                    violations=['图片下载失败'],
                    confidence=0.0
                )
            
            # 基础检查
            basic_violations = self._basic_image_check(image)
            
            # 内容检查（简化版本，实际应用中可以集成更复杂的AI模型）
            content_violations = self._content_check(image)
            
            # 合并违规结果
            all_violations = basic_violations + content_violations
            
            # 计算置信度
            confidence = self._calculate_confidence(all_violations, image)
            
            result = ImageAuditResult(
                is_safe=len(all_violations) == 0,
                violations=all_violations,
                confidence=confidence
            )
            
            if all_violations:
                audit_logger.info(f"图片审核发现违规: {all_violations}")
            
            return result
            
        except Exception as e:
            audit_logger.error(f"图片审核失败: {e}")
            return ImageAuditResult(
                is_safe=False,
                violations=['审核异常'],
                confidence=0.0
            )
    
    def audit_images(self, image_urls: List[str]) -> List[ImageAuditResult]:
        """批量审核图片"""
        results = []
        for url in image_urls:
            result = self.audit_image(url)
            results.append(result)
        return results
    
    def _download_image(self, image_url: str) -> Image.Image:
        """下载图片"""
        try:
            # 设置请求头
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            
            # 下载图片
            response = requests.get(image_url, headers=headers, timeout=10)
            response.raise_for_status()

            # 打开图片
            from io import BytesIO
            image = Image.open(BytesIO(response.content))
            return image
            
        except Exception as e:
            audit_logger.error(f"下载图片失败 {image_url}: {e}")
            return None
    
    def _basic_image_check(self, image: Image.Image) -> List[str]:
        """基础图片检查"""
        violations = []
        
        try:
            # 检查图片尺寸
            width, height = image.size
            
            # 图片太小可能是像素图或无效图片
            if width < 50 or height < 50:
                violations.append('图片尺寸过小')
            
            # 图片太大可能包含过多信息
            if width > 5000 or height > 5000:
                violations.append('图片尺寸过大')
            
            # 检查图片格式
            if image.format not in ['JPEG', 'PNG', 'GIF', 'WEBP']:
                violations.append('不支持的图片格式')
            
            # 检查图片模式
            if image.mode not in ['RGB', 'RGBA', 'L']:
                violations.append('异常图片模式')
                
        except Exception as e:
            audit_logger.error(f"基础图片检查失败: {e}")
            violations.append('图片格式异常')
        
        return violations
    
    def _content_check(self, image: Image.Image) -> List[str]:
        """内容检查（简化版本）"""
        violations = []
        
        try:
            # 转换为numpy数组进行分析
            img_array = np.array(image)
            
            # 检查图片亮度分布
            if len(img_array.shape) == 3:
                # 彩色图片
                brightness = np.mean(img_array)
                
                # 过暗的图片可能包含不当内容
                if brightness < 30:
                    violations.append('图片过暗')
                
                # 检查颜色分布
                color_violations = self._check_color_distribution(img_array)
                violations.extend(color_violations)
            
            # 检查图片复杂度
            complexity_violations = self._check_image_complexity(img_array)
            violations.extend(complexity_violations)
            
        except Exception as e:
            audit_logger.error(f"图片内容检查失败: {e}")
            violations.append('内容分析异常')
        
        return violations
    
    def _check_color_distribution(self, img_array: np.ndarray) -> List[str]:
        """检查颜色分布"""
        violations = []
        
        try:
            # 计算各颜色通道的标准差
            if len(img_array.shape) == 3 and img_array.shape[2] >= 3:
                r_std = np.std(img_array[:, :, 0])
                g_std = np.std(img_array[:, :, 1])
                b_std = np.std(img_array[:, :, 2])
                
                # 颜色分布过于单一可能是纯色图片或异常图片
                if r_std < 10 and g_std < 10 and b_std < 10:
                    violations.append('颜色分布异常')
                
                # 检查是否存在大量红色（可能是血腥内容）
                red_ratio = np.mean(img_array[:, :, 0]) / 255.0
                if red_ratio > 0.7:
                    violations.append('疑似血腥内容')
                    
        except Exception as e:
            audit_logger.error(f"颜色分布检查失败: {e}")
        
        return violations
    
    def _check_image_complexity(self, img_array: np.ndarray) -> List[str]:
        """检查图片复杂度"""
        violations = []
        
        try:
            # 转换为灰度图
            if len(img_array.shape) == 3:
                gray = np.mean(img_array, axis=2)
            else:
                gray = img_array
            
            # 计算梯度来评估图片复杂度
            grad_x = np.gradient(gray, axis=1)
            grad_y = np.gradient(gray, axis=0)
            gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
            
            # 计算平均梯度
            avg_gradient = np.mean(gradient_magnitude)
            
            # 梯度过低可能是纯色图片或低质量图片
            if avg_gradient < 5:
                violations.append('图片质量过低')
            
            # 梯度过高可能包含过多细节（如二维码等）
            if avg_gradient > 50:
                violations.append('疑似包含二维码或复杂图案')
                
        except Exception as e:
            audit_logger.error(f"图片复杂度检查失败: {e}")
        
        return violations
    
    def _calculate_confidence(self, violations: List[str], image: Image.Image) -> float:
        """计算审核置信度"""
        if not violations:
            return 0.9  # 没有违规，高置信度
        
        # 基础置信度
        base_confidence = 0.7
        
        # 根据违规数量调整
        violation_penalty = len(violations) * 0.1
        
        # 根据图片质量调整
        try:
            width, height = image.size
            size_factor = min(width * height / (1000 * 1000), 1.0)  # 归一化到[0,1]
            quality_bonus = size_factor * 0.1
        except:
            quality_bonus = 0
        
        confidence = base_confidence - violation_penalty + quality_bonus
        return max(0.1, min(confidence, 1.0))


# 全局图片审核器实例
image_auditor = ImageAuditor()
